from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
from numpy import array, random, linspace, pi, ravel, cos, sin, empty, mgrid, arange
from mayavi import mlab
import networkx as nx
from tvtk.api import tvtk
from mayavi.sources.builtin_surface import BuiltinSurface
import csv
from scipy.special import sph_harm

class Form(HasTraits):
    scene = Instance(MlabSceneModel, ())
    file = File('temperature.axdt')

    button = Button('显示温度场')

    def _button_fired(self):
        with open(self.file) as file:
            fileStrings = file.readlines()
        a = fileStrings.index('[Data]\n')
        b = fileStrings.index('[Faces]\n')
        c1 = []
        for value in fileStrings[2:b]:
            if len(tmp := (value.strip().split(','))) == 4:
                c1.append(tmp)
        vt = np.array(c1).astype(float)
        v = vt[:, :3]
        t = vt[:, 3]
        c1 = []
        for value in fileStrings[b + 1:]:
            if len(tmp := (value.strip().split(','))) == 3:
                c1.append(tmp)
        f = np.array(c1).astype(int)
        self.v = v
        self.f = f
        self.t = t
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.clf()
        surf = mlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=t)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0))
        mlab.colorbar()

    buttonIR = Button('计算红外辐射')

    def _buttonIR_fired(self):
        v = self.v
        f = self.f
        t = self.t
        thetas = np.linspace(-40, 40, 41)
        phis = np.linspace(0, 360, 181)
        ir = np.zeros((thetas.shape[0], phis.shape[0]))
        v0 = v[f[:, 0], :]
        v1 = v[f[:, 1], :]
        v2 = v[f[:, 2], :]
        tf = (t[f[:, 0]] + t[f[:, 1]] + t[f[:, 2]]) / 3
        n = np.cross(v1 - v0, v2 - v0)
        a = np.sqrt((n * n).sum(1)) / 2
        n /= 2 * a[:, np.newaxis]
        for tt, theta in enumerate(thetas):
            costheta = np.cos(theta * np.pi / 180)
            sintheta = np.sin(theta * np.pi / 180)
            print(tt)
            for pp, phi in enumerate(phis):
                cosphi = np.cos(phi * np.pi / 180)
                sinphi = np.sin(phi * np.pi / 180)
                s = np.array([costheta * cosphi, costheta * sinphi, sintheta])
                ns = n @ s
                light = ns < -1e-6
                ir[tt, pp] = np.sum(5.67e-8 * tf[light] ** 4 / np.pi * (-ns[light]) * a[light])
        self.ir = ir
        self.thetas = thetas
        self.phis = phis
        print(ir.max())

    buttonShowIR = Button('显示红外辐射')

    def _buttonShowIR_fired(self):
        thetas, phis = np.meshgrid(self.thetas, self.phis, indexing='ij')
        costhetas = np.cos(thetas * np.pi / 180)
        sinthetas = np.sin(thetas * np.pi / 180)
        cosphis = np.cos(phis * np.pi / 180)
        sinphis = np.sin(phis * np.pi / 180)
        x = costhetas * cosphis
        y = costhetas * sinphis
        z = sinthetas
        ir = self.ir * 5
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.mesh(ir * x, ir * y, ir * z, opacity=.5)
        mlab.colorbar()

    buttonShowIRSphere = Button('显示红外辐射球面')

    def _buttonShowIRSphere_fired(self):
        thetas, phis = np.meshgrid(self.thetas, self.phis, indexing='ij')
        costhetas = np.cos(thetas * np.pi / 180)
        sinthetas = np.sin(thetas * np.pi / 180)
        cosphis = np.cos(phis * np.pi / 180)
        sinphis = np.sin(phis * np.pi / 180)
        x = costhetas * cosphis
        y = costhetas * sinphis
        z = sinthetas
        ir = self.ir * 5
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.mesh(x, y, z, scalars=ir, opacity=.5)
        mlab.colorbar()

    buttonClear = Button('清除视景')

    def _buttonClear_fired(self):
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.clf()

    buttonShowIRScene = Button('测试')

    def _buttonShowIRScene_fired(self):
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        surf = mlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest1 = Button('测试1')

    def _buttonTest1_fired(self):
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        def polydata():
            # The numpy array data.
            points = np.array([[0, -0.5, 0], [1.5, 0, 0], [0, 1, 0], [0, 0, 0.5], [-1, -1.5, 0.1], [0, -1, 0.5], [-1, -0.5, 0], [1, 0.8, 0]], 'f')
            triangles = np.array([[0, 1, 3], [1, 2, 3], [1, 0, 5], [2, 3, 4], [3, 0, 4], [0, 5, 4], [2, 4, 6], [2, 1, 7]])
            scalars = np.random.random(points.shape)
            # The TVTK dataset.
            mesh = tvtk.PolyData(points=points, polys=triangles)
            mesh.point_data.scalars = scalars
            mesh.point_data.scalars.name = 'scalars'
            return mesh

        def view(dataset):
            surf = mlab.pipeline.surface(dataset, opacity=0.1)
            mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(polydata())

    buttonTest2 = Button('测试2')

    def _buttonTest2_fired(self):
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # scalars = np.array([[0, 1, 0]])
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.point_data.scalars = scalars
        # mesh.cell_data.scalars = scalars
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest3 = Button('测试3')

    def _buttonTest3_fired(self):
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([[0, .5, 1]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.cell_data.scalars = scalars
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest4 = Button('测试4')

    def _buttonTest4_fired(self):
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        surf = mlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    buttonTest5 = Button('测试5，默认红外文件')

    def _buttonTest5_fired(self):

        with open('temperature.axdt') as file:
            fileStrings = file.readlines()
        a = fileStrings.index('[Data]\n')
        b = fileStrings.index('[Faces]\n')
        c1 = []
        for value in fileStrings[2:b]:
            if len(tmp := (value.strip().split(','))) == 4:
                c1.append(tmp)
        vt = np.array(c1).astype(float)
        v = vt[:, :3]
        t = vt[:, 3]
        c1 = []
        for value in fileStrings[b + 1:]:
            if len(tmp := (value.strip().split(','))) == 3:
                c1.append(tmp)
        f = np.array(c1).astype(int)
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        surf = mlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=t)
        mlab.colorbar()
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )
        pass

    v = None
    f = None
    t = None

    view = View(
        UItem('scene', editor=SceneEditor(scene_class=MayaviScene)),
        HGroup(
            UItem('button'),
            UItem('buttonIR'),
            UItem('buttonClear'),
            UItem('buttonShowIR'),
            UItem('buttonShowIRScene'),
            UItem('buttonShowIRSphere'),
            UItem('buttonTest1'),
            UItem('buttonTest2'),
            UItem('buttonTest3'),
            UItem('buttonTest4'),
            UItem('buttonTest5'),
            UItem('file'),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )

#
# from PyQt5.QtWidgets import *
#
#
# class FormWidget(QWidget):
#     def __init__(self):
#         QWidget.__init__(self)
#         layout = QHBoxLayout(self)
#         layout.setContentsMargins(0, 0, 0, 0)
#         layout.setSpacing(0)
#         self.form = Form()
#         layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)
#

# formWidget = FormWidget()
if __name__ == '__main__':
    Form().configure_traits()
    # import sys
    # app = QApplication(sys.argv)
    # formWidget.show()
    # sys.exit(app.exec_())



